import threading
import time
from typing import Optional, List, Tuple
from dataclasses import dataclass
import heapq
import io

from loguru import logger


@dataclass
class Chunk:
    start: int
    end: int
    data: bytes

    def __lt__(self, other):
        return self.start < other.start
    
    def __len__(self):
        return self.end - self.start


class ChunkIOError(Exception):
    """Base exception for ChunkIO errors"""

    pass


class DataOverlapError(ChunkIOError):
    """Raised when overlapping data chunks are detected"""

    pass


class TimeoutError(ChunkIOError):
    """Raised when read operation times out"""

    pass


class EOFError(ChunkIOError):
    """Raised when EOF command has incorrect parameters"""

    pass


class ChunkIO:
    def __init__(self):
        """Initialize ChunkIO with empty state"""
        self._chunks: List[Chunk] = []  # Will be used as a heap
        self._read_position = 0
        self._is_eof = False
        self._has_error = False
        self._lock = threading.Lock()
        self._data_available = threading.Event()

    def write(self, start_bytes: int, end_bytes: int, data: Optional[bytes]) -> None:
        """
        Write data chunk to the buffer. Thread-safe operation.

        Args:
            start_bytes (int): Start position of the chunk
            end_bytes (int): End position of the chunk
            data (Optional[bytes]): Data to write. None indicates EOF

        Raises:
            ChunkIOError: If error flag is set
            EOFError: If EOF command has incorrect parameters
            DataOverlapError: If chunk overlaps with existing data
        """
        if self._has_error:
            raise ChunkIOError("Cannot write: error flag is set")

        with self._lock:
            if data is None:  # EOF command
                # logger.info(f"EOF command: {start_bytes} {end_bytes}")
                if start_bytes != end_bytes:
                    raise EOFError(
                        "EOF command must have start_bytes and end_bytes equal to file length"
                    )
                heapq.heappush(self._chunks, Chunk(start_bytes, end_bytes, None))
                # logger.info(f"Pushed EOF chunk: {start_bytes} {end_bytes}")
                self._data_available.set()
                return
            # logger.info(f"Writing {start_bytes} {end_bytes}")
            assert start_bytes <= end_bytes
            assert start_bytes >= 0
            assert data is not None
            assert len(data) == end_bytes - start_bytes
            # logger.info(f"Pushing chunk: {start_bytes} {end_bytes}")
            heapq.heappush(self._chunks, Chunk(start_bytes, end_bytes, data))

            self._data_available.set()

    def read(self, n: int, timeout: float = 5.0) -> Optional[bytes]:
        """
        Read n bytes from the buffer. Thread-safe operation.

        Args:
            n (int): Number of bytes to read
            timeout (float): Maximum time in seconds to wait for data before raising TimeoutError

        Returns:
            Optional[bytes]: Read data or None if EOF reached

        Raises:
            ChunkIOError: If error flag is set
            TimeoutError: If timeout is reached while waiting for data
        """
        # logger.info(f"Reading {n} bytes from ChunkIO")
        if self._has_error:
            logger.error(f"Cannot read: error flag is set")
            raise ChunkIOError("Cannot read: error flag is set")

        if n <= 0:
            raise ValueError("n must be positive")

        if self._is_eof:
            # logger.info(f"EOF reached in ChunkIO")
            return b''

        start_time = time.time()

        while True:
            with self._lock:
                result, EOF = self._get_available_data(n)
                self._is_eof = EOF
                if result is not None:
                    assert isinstance(result, bytes), f"Result is not bytes: {type(result)}"
                    self._read_position += len(result)
                    # logger.info(f"Read {len(result)} bytes from ChunkIO")
                    return result

                # Wait for more data outside the lock
                self._data_available.clear()
            max_wait_time = timeout - (time.time() - start_time)
            if max_wait_time <= 0:
                logger.error(f"Timeout waiting for data after {timeout} seconds")
                raise TimeoutError(f"Timeout waiting for data after {timeout} seconds")
            if not self._data_available.wait(timeout=max_wait_time):
                logger.error(f"Timeout waiting for data after {timeout} seconds")
                raise TimeoutError(f"Timeout waiting for data after {timeout} seconds")
        raise RuntimeError("Should not reach here")

    def _get_available_data(self, n: int) -> Tuple[Optional[bytes], bool]:
        """
        Get available data from the buffer.

        Args:
            n (int): Number of bytes to read

        Returns:
            Optional[bytes]: Available data or None if not enough data available
            bool: True if EOF reached
        """
        # logger.info(f"Getting available data: {n}")
        head_chunk = Chunk(self._read_position, self._read_position, bytearray())
        EOF = False
        while self._chunks:
            # check if the next chunk is contiguous
            if self._chunks[0].start > head_chunk.end:
                # logger.info(f"Chunk {self._chunks[0].start}~{self._chunks[0].end} is not contiguous with head_chunk {head_chunk.start}~{head_chunk.end}")
                break
            # merge the two chunks
            assert self._chunks[0].start == head_chunk.end
            if self._chunks[0].data is None:
                assert self._chunks[0].end == self._chunks[0].start
                EOF = True
                break
            head_chunk.end = self._chunks[0].end
            head_chunk.data.extend(self._chunks[0].data)
            heapq.heappop(self._chunks)
        # logger.info(f"head_chunk: {head_chunk.start}~{head_chunk.end} {EOF} {self._read_position}")
        # logger.info(f"chunks: {self._chunks}")
        if len(head_chunk) < n:
            if EOF:
                return bytes(head_chunk.data), EOF
            heapq.heappush(self._chunks, head_chunk)
            return None, False
        elif len(head_chunk) == n:
            return bytes(head_chunk.data), EOF
        else:
            return_data = head_chunk.data[:n]
            head_chunk.data = head_chunk.data[n:]
            head_chunk.start = self._read_position + n
            heapq.heappush(self._chunks, head_chunk)
            return bytes(return_data), False